/*
 * Main.java
 *
 * Created on Dec 1, 2011
 *
 * Copyright (c) 2011 Artur Rataj
 *
 * This code is distributed under the terms of the GNU Library
 * General Public License, either version 3 of the license or, at
 * your option, any later version.
 */

package pl.gliwice.iitis.spatialnn;

import java.io.*;
import java.util.*;
import java.awt.image.BufferedImage;
import java.awt.image.Raster;
import java.awt.image.WritableRaster;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.imageio.ImageIO;

import pl.gliwice.iitis.graphutils.Point;
import pl.gliwice.iitis.graphutils.view.DrawMode;
import pl.gliwice.iitis.spatialnn.network.*;
import pl.gliwice.iitis.spatialnn.network.construct.*;
import pl.gliwice.iitis.spatialnn.network.io.ParseException;
import pl.gliwice.iitis.spatialnn.network.io.parser.NetworkIO;
import pl.gliwice.iitis.spatialnn.network.view.NetworkView;
import pl.gliwice.iitis.spatialnn.propagate.*;
import pl.gliwice.iitis.spatialnn.propagate.image.PropagateImage;

/**
 * The main class. Parses command--line arguments.
 * 
 * @author Artur Rataj
 */
public class Main {
    protected enum Task {
        IO_TEST,
        DEPTH,
        FREQUENCY,
        PROPAGATION,
        BLUR,
    }
    public static final boolean TIME_MEASURE = false;
    /**
     * Constructs a test network.
     * 
     * @param task type of task
     * @param network network to add components to, prune and pack
     */
    protected static void construct(Task task, Network network) {
        Wood w = new Wood();
        NetworkStats stats = new NetworkStats();
        stats.membrane = new MembraneStats(
                network.getOutMembrane().type.getConstructionParams());
        switch(task) {
            case DEPTH:
                stats.addLayer(NeuronType.MINP, -1, false);
                stats.addLayer(NeuronType.MEAN, 0, false);
                stats.addLayer(NeuronType.MEAN, 1, false);
//                stats.addLayer(NeuronType.MEAN, 2, false);
//                stats.addLayer(NeuronType.MEAN, 3, false);
                stats.addLayer(NeuronType.ABS, 2, false);
                break;

            case FREQUENCY:
                stats.addLayer(NeuronType.MEAN, -1, true);
                if(false) {
                    final int N = 16;
                    stats.setTurn(0, N, 2, 1);
                    for(int i = 0; i < N; ++i) {
                        stats.addLayer(NeuronType.ABS, i*2 + 0, false);
                        stats.addLayer(NeuronType.MEAN, i*2 + 1, false);
                    }
                } else {
                    stats.setTurn(0, 1, 1, 0);
                    stats.addLayer(NeuronType.ABS, 0, false);
                }
                break;

            case BLUR:
                stats.addLayer(NeuronType.MEAN, -1, false);
                stats.addLayer(NeuronType.MEAN, 0, false);
                stats.addLayer(NeuronType.MEAN, 1, false);
                stats.addLayer(NeuronType.MEAN, 2, false);
                break;

            default:
                throw new RuntimeException("unsupported task");
        }
        w.construct(network, stats);
    }
    /**
     * Propagates signals through a test network.
     * 
     * @param network network
     */
    public static void propagate(Network network) {
        network.membraneMap.get("left").setRes(5);
        network.membraneMap.get("right").setRes(5);
        network.membraneMap.get("out").setRes(5);
        Propagate p = new Propagate(network);
        Cache c = p.getCache();
        Map<String, double[]> in = new HashMap<>();
        double[] inLeft = {
            0.0, 1.0, 2.0, 3.0, 4.0,
        };
        double[] inRight = {
            -1, -2, -3, -4, -5,
        };
        in.put("left", inLeft);
        in.put("right", inRight);
        double[] out = p.get(in, c);
        double[] outTest = {
             0.3, 0.3, 1.15, 2.0, 2.0,
        };
        if(out.length != outTest.length)
            throw new RuntimeException("invalid output size");
        for(int i = 0; i < outTest.length; ++i)
            if(out[i] != outTest[i])
                throw new RuntimeException("invalid output value");
    }
    /**
     * Propagates an image to blur it.
     * 
     * @param network network
     */
    public static void propagateBlur(Network network) {
        final int RES = 100;
        network.membraneMap.get("input").setRes(RES);
        network.membraneMap.get("output").setRes(RES);
        PropagateImage p = new PropagateImage(network);
        Cache c = p.getCache();
        Map<String, BufferedImage> in = new HashMap<>();
        try {
            BufferedImage input = ImageIO.read(new File("images/Midd1/view1.png"));
            in.put("input", input);
            p.prepareImages(in, 0.2);
            BufferedImage out = p.get(in, 1.0, c);
            ImageIO.write(out, "PNG", new File("/tmp/imageblur.png"));
        } catch (IOException e) {
            throw new RuntimeException("io error: " + e.toString());
        }
    }
    /**
     * Propagates an image to get spectrum.
     * 
     * @param network network
     */
    public static void propagateFrequency(Network network) {
        final int RES = 100;
        network.membraneMap.get("space").setRes(RES);
        network.membraneMap.get("feature").setRes(RES);
        PropagateImage p = new PropagateImage(network);
        Cache c = p.getCache();
        Map<String, BufferedImage> in = new HashMap<>();
        try {
            BufferedImage input = ImageIO.read(
                    new File("bars-angled.png")
                    //new File("freq_wood.png")
                    //new File("usc_mosaic_2.png")
                    );
            in.put("space", input);
            p.prepareImages(in, 1.0);
            BufferedImage out = p.get(in, 1, c);
            ImageIO.write(out, "PNG", new File("/tmp/imagefrequency.png"));
        } catch (IOException e) {
            throw new RuntimeException("io error: " + e.toString());
        }
    }
    protected static int clamp(int v, int min, int max) {
        if(v < min)
            v = min;
        else if(v > max)
            v = max;
        return v;
    }
    /**
     * Estimates MSE on basis of ground truth and depth map.
     * 
     * @param tLeft left truth
     * @param tRight right truth
     * @param out depth map estimation from a network
     * @param b b in bx + a for brightness to disparity
     * @param a a in bx + a for brightness to disparity
     */
    protected static double mseDepth(Raster tLeft, Raster tRight, BufferedImage out,
            double b, double a, double brightnessScale, double yTop) {
        final double GAMMA = 1.0;
        double mse = 0.0;
        int w = tLeft.getWidth();
        int h = tLeft.getHeight();
        if(w != tRight.getWidth())
            throw new RuntimeException("width mismatch");
        if(h != tRight.getHeight())
            throw new RuntimeException("height mismatch");
        if(w != out.getWidth() || h != out.getHeight())
            out = PropagateImage.resize(out, w, h);
        Raster tEst = out.getRaster();
        BufferedImage control = new BufferedImage(w, h, BufferedImage.TYPE_BYTE_GRAY);
        WritableRaster r = control.getRaster();
        double count  = 0;
        double layers = 0.02*Math.cos(Math.PI/4.0)*2;
        double c1Sum = 0.0;
        PriorityQueue<Double> queue = new PriorityQueue<>();
        for(int y = 0; y < h; ++y)
            for(int x = 0; x < w; ++x) {
                int cLeft = tLeft.getSample(x, y, 0);
                int cRight = tLeft.getSample(x, y, 0);
                if(cLeft != 0 && cRight != 0) {
                    double c1 = ((cLeft + cRight + 0.5)/2 - 130)/3.0;
                    c1Sum += c1;
                    double c2 = 0;
                    final int W = 0;
                    for(int q = -W; q <= W; ++q)
                        for(int p = -W; p <= W; ++p) {
                            c2 += (yTop - layers - tEst.getSample(
                                    clamp(x + p, 0, w - 1), clamp(y + q, 0, h - 1),
                                    0)*brightnessScale)*b + a;
                        }
                    c2 = c2/((2*W + 1)*(2*W + 1));
                    // c2 = -1.725;
                    double c = c2*3 + 130;
                    if(c > 0xff)
                        c = 0xff;
                    if(c < 10)
                        c = 10;
                    r.setSample(x, y, 0, c);
                    double component = Math.pow(c2 - c1, 2.0);
                    if(GAMMA == 1.0)
                        mse += component;
                    else {
                        queue.add(component);
                    }
                    ++count;
                }
            }
        System.out.println("mean truth = " + (c1Sum/count));
        if(GAMMA == 1.0)
            mse = mse/count;
        else {
            int maxCount = (int)Math.round(count*GAMMA);
            for(int i = 0; i < maxCount; ++i)
                mse += queue.poll();
            mse = mse/maxCount;
        }
        try {
            ImageIO.write(control, "PNG", new File("/tmp/control.png"));
        } catch (IOException e) {
            throw new RuntimeException("io: " + e.toString());
        }
        return mse;
    }
    /**
     * Propagates stereo images through a depth map network.
     * 
     * @param network network
     */
    public static void propagateDepth(Network network) {
        final int RES = 100;
        network.membraneMap.get("left").setRes(RES);
        network.membraneMap.get("right").setRes(RES);
        network.membraneMap.get("depth").setRes(RES);
        PropagateImage p = new PropagateImage(network);
        Cache c = p.getCache();
        Map<String, BufferedImage> in = new HashMap<>();
        try {
            //BufferedImage left = ImageIO.read(new File("images/Bowling2/view1.png"));
            //BufferedImage right = ImageIO.read(new File("images/Bowling2/view5.png"));
            //BufferedImage left = ImageIO.read(new File("left_bar.png"));
            //BufferedImage right = ImageIO.read(new File("right_bar.png"));
            BufferedImage left = ImageIO.read(new File("left_wood.png"));
            BufferedImage right = ImageIO.read(new File("right_wood.png"));
            in.put("left", left);
            in.put("right", right);
            p.prepareImages(in, 0.3);
            double brightnessScale = 0.5;
            BufferedImage out = p.get(in, brightnessScale, c);
            ImageIO.write(out, "PNG", new File("/tmp/imagedepth.png"));
            BufferedImage tLeft = ImageIO.read(new File("doc/depth-wood-left.png"));
            BufferedImage tRight = ImageIO.read(new File("doc/depth-wood-right.png"));
            Membrane mLeft = network.membraneMap.get("left");
            Membrane mRight = network.membraneMap.get("right");
            double yTop = mLeft.locBeg.y;
            double xOffset = (mRight.locEnd.x + mRight.locBeg.x)/2.0 -
                    (mLeft.locEnd.x + mLeft.locBeg.x)/2.0;
            double xExtend = (mRight.locEnd.x - mRight.locBeg.x)/2.0 +
                    (mLeft.locEnd.x - mLeft.locBeg.x)/2.0;
            // (mRight.locEnd.x - mRight.locBeg.x)/(mLeft.locEnd.x - mLeft.locBeg.x)
            double spread = network.symmetricInputSpread;
            double res = left.getWidth();
            double b = -2*Math.tan(spread)*res/xExtend;
            double a = +xOffset*res/xExtend;
            final int MAX_SAMPLE = (1 << 16) - 1;
            double mse = mseDepth(tLeft.getRaster(), tRight.getRaster(), out, b, a,
                    1.0/(brightnessScale*MAX_SAMPLE), yTop);
            System.out.println("MSE = " + mse);
        } catch (IOException e) {
            throw new RuntimeException("io error: " + e.toString());
        }
    }
    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) {
        //Task task = Task.IO_TEST;
        Task task = Task.DEPTH;
        //Task task = Task.FREQUENCY;
        //Task task = Task.PROPAGATION;
        //Task task = Task.BLUR;
        String[] args_ = new String[1];
        // if the network needs additional construction
        boolean construct;
        switch(task) {
            case IO_TEST:
                 args_[0] = "test.network";
                 construct = false;
                 break;
                
            case DEPTH:
                 args_[0] = "straight.network";
                 construct = true;
                 break;
                
            case FREQUENCY:
                args_[0] = "turn.network";
                construct = true;
                break;
                        
            case PROPAGATION:
                args_[0] = "test_propagation.network";
                construct = false;
                break;
                        
            case BLUR:
                args_[0] = "siso.network";
                construct = true;
                break;
                        
            default:
                throw new RuntimeException("unknown task");
        }
        args = args_;
        if(args.length != 1)
            System.out.println("syntax: <network file>");
        else {
            Network network = null;
            String fileName = args[0];
            long start = System.currentTimeMillis();
            try {
                network = NetworkIO.read(fileName);
            } catch(ParseException e) {
                throw new RuntimeException("parse error: " + e.toString());
            }
            if(construct && network != null)
                construct(task, network);
            if(!TIME_MEASURE) {
                DrawMode dm = new DrawMode(70, true);
                NetworkView pv = new NetworkView(dm, 1000, 0.5f,
                        false, true);
                BufferedImage image = pv.render(network);
                try {
                    ImageIO.write(image, "PNG", new File("/tmp/network.png"));
                    NetworkIO.write(network, "/tmp/a.network");
                } catch(IOException e) {
                    System.out.println("write error: " + e.toString());
                }
            }
            switch(task) {
                case DEPTH:
                    propagateDepth(network);
                    break;

                case PROPAGATION:
                    propagate(network);
                    break;
                    
                case FREQUENCY:
                    propagateFrequency(network);
                    break;
                    
                case IO_TEST:
                    break;
                
                case BLUR:
                    propagateBlur(network);
                    break;

                default:
                    throw new RuntimeException("unknown task");
            }
            long end = System.currentTimeMillis();
            if(TIME_MEASURE || true)
                System.out.println("Joint construction and propagation time was "+(end-start)/1000.0+" s.");
        }
    }
}
